import struct
from enum import Enum
from decimal import Decimal
from typing import Dict
from abc import ABC
from typing import Any, List, Union


class BaseEvent(ABC):
    """Base class for all event types."""

    def __init__(self, data: Dict):
        if not isinstance(data, dict):
            raise TypeError("Input data must be dict.")
        self._origin_data = data


class MindSporeOpIndexEnum(Enum):
    THREAD_ID = 0
    FLOW_ID = 1
    STEP = 2
    START_TIME_NS = 3
    END_TIME_NS = 4
    PROCESS_ID = 5
    MODULE_INDEX = 6
    EVENT_INDEX = 7
    STAGE_INDEX = 8
    LEVEL = 9
    IS_ASYNC = 10
    IS_GRAPH_DATA = 11
    IS_STAGE = 12
    NAME = 13


class MindSporeOpEvent(BaseEvent):
    FIX_DATA_FORMAT = "<5Qi3Hb3?"
    FIX_DATA_SIZE = struct.calcsize(FIX_DATA_FORMAT)

    def __init__(self, data: Dict):
        super().__init__(data)
        self.fix_size_data = struct.unpack(
            self.FIX_DATA_FORMAT, self._origin_data.get("fix_size_bytes")
        )

    @property
    def thread_id(self):
        return self.fix_size_data[MindSporeOpIndexEnum.THREAD_ID.value]

    @property
    def flow_id(self):
        return self.fix_size_data[MindSporeOpIndexEnum.FLOW_ID.value]

    @property
    def step(self):
        return self.fix_size_data[MindSporeOpIndexEnum.STEP.value]

    @property
    def start_time_ns(self):
        return self.fix_size_data[MindSporeOpIndexEnum.START_TIME_NS.value]

    @property
    def end_time_ns(self):
        return self.fix_size_data[MindSporeOpIndexEnum.END_TIME_NS.value]

    @property
    def process_id(self):
        return self.fix_size_data[MindSporeOpIndexEnum.PROCESS_ID.value]

    @property
    def module_index(self):
        return self.fix_size_data[MindSporeOpIndexEnum.MODULE_INDEX.value]

    @property
    def event_index(self):
        return self.fix_size_data[MindSporeOpIndexEnum.EVENT_INDEX.value]

    @property
    def stage_index(self):
        return self.fix_size_data[MindSporeOpIndexEnum.STAGE_INDEX.value]

    @property
    def level(self):
        return self.fix_size_data[MindSporeOpIndexEnum.LEVEL.value]

    @property
    def is_async(self):
        return self.fix_size_data[MindSporeOpIndexEnum.IS_ASYNC.value]

    @property
    def is_graph_data(self):
        return self.fix_size_data[MindSporeOpIndexEnum.IS_GRAPH_DATA.value]

    @property
    def is_stage(self):
        return self.fix_size_data[MindSporeOpIndexEnum.IS_STAGE.value]

    @property
    def name(self):
        return self._origin_data.get(MindSporeOpIndexEnum.NAME.value, "")

    def __str__(self):
        return f"MindSporeOpEvent(name={self.name}, thread_id={self.thread_id}, flow_id={self.flow_id}, step={self.step}, start_time_ns={self.start_time_ns}, end_time_ns={self.end_time_ns}, process_id={self.process_id}, module_index={self.module_index}, event_index={self.event_index}, stage_index={self.stage_index}, level={self.level}, is_async={self.is_async}, is_graph_data={self.is_graph_data}, is_stage={self.is_stage})"


class TLVDecoder:
    """
    The class of TLV format data decoder.

    Args:
        all_bytes(bytes): all the bytes data of tlv format binary file.
        event_class(MindSporeOpEvent): class of events to decode.
        fix_data_struct_size(int): a constant value.

    Return:
        List[MindSporeOpEvent]: event list of input event class.
    """
    _type_len = 2
    _length_len = 4

    @classmethod
    def decode(cls, all_bytes: bytes, event_class: MindSporeOpEvent, fix_data_struct_size: int):
        """Decode all the data."""
        result_data = []
        records = cls.tlv_list_decode(all_bytes)
        for record in records:
            if fix_data_struct_size > len(record):
                print("The collected data has been lost 1")
                continue
            fix_data_bytes = record[0: fix_data_struct_size]
            tlv_fields = cls.tlv_list_decode(record[fix_data_struct_size:], is_field=True)
            tlv_fields["fix_size_bytes"] = fix_data_bytes
            result_data.append(event_class(tlv_fields))
        return result_data

    @classmethod
    def tlv_list_decode(cls, tlv_bytes: bytes, is_field: bool = False) -> Union[Dict, List]:
        """Decode TLV format data."""
        result_data = {} if is_field else []
        index = 0
        all_bytes_len = len(tlv_bytes)
        while index < all_bytes_len:
            if index + cls._type_len > all_bytes_len:
                print("The collected data has been lost 2")
                break
            type_id = struct.unpack("<H", tlv_bytes[index: index + cls._type_len])[0]
            index += cls._type_len
            if index + cls._length_len > all_bytes_len:
                print("The collected data has been lost 2")
                break
            value_len = struct.unpack("<I", tlv_bytes[index: index + cls._length_len])[0]
            index += cls._length_len
            if index + value_len > all_bytes_len:
                print("The collected data has been lost 3")
                break
            value = tlv_bytes[index: index + value_len]
            index += value_len
            if is_field:
                try:
                    result_data[type_id] = bytes.decode(value)
                except UnicodeDecodeError:
                    print(f"The collected data can't decode by bytes.decode: {value}")
                    result_data[type_id] = 'N/A'
            else:
                result_data.append(value)
        return result_data

if __name__ == "__main__":
    with open("./tmp/mindspore.op_range", "rb") as f:
        all_bytes = f.read()

    print(MindSporeOpEvent.FIX_DATA_SIZE)
    op_range_data = TLVDecoder.decode(
        all_bytes, MindSporeOpEvent, MindSporeOpEvent.FIX_DATA_SIZE
    )
    print(
        f"op_range_data len : {len(op_range_data)}, op_range_data: {op_range_data[0]}"
    )
    print(op_range_data[0]._origin_data)
